{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "931f7017-d2e2-4b1a-8fdf-9ddf119fc8ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries of pydantic and the llama-cpp-agent framework.\n",
    "from enum import Enum\n",
    "from typing import List\n",
    "\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "from llama_cpp_agent import LlamaCppAgent\n",
    "\n",
    "from llama_cpp import Llama\n",
    "from llama_cpp_agent.providers import LlamaCppPythonProvider\n",
    "\n",
    "\n",
    "\n",
    "from llama_cpp_agent.llm_output_settings import LlmStructuredOutputSettings, LlmStructuredOutputType\n",
    "from llama_cpp_agent.providers import TGIServerProvider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "059444fc-6f34-4403-bea2-d19b6457b792",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_model_loader: loaded meta data with 27 key-value pairs and 339 tensors from ./deepseek-r1-distill-qwen-7b-q4_k_m.gguf (version GGUF V3 (latest))\n",
      "llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n",
      "llama_model_loader: - kv   0:                       general.architecture str              = qwen2\n",
      "llama_model_loader: - kv   1:                               general.type str              = model\n",
      "llama_model_loader: - kv   2:                               general.name str              = DeepSeek R1 Distill Qwen 7B\n",
      "llama_model_loader: - kv   3:                           general.basename str              = DeepSeek-R1-Distill-Qwen\n",
      "llama_model_loader: - kv   4:                         general.size_label str              = 7B\n",
      "llama_model_loader: - kv   5:                            general.license str              = mit\n",
      "llama_model_loader: - kv   6:                          qwen2.block_count u32              = 28\n",
      "llama_model_loader: - kv   7:                       qwen2.context_length u32              = 131072\n",
      "llama_model_loader: - kv   8:                     qwen2.embedding_length u32              = 3584\n",
      "llama_model_loader: - kv   9:                  qwen2.feed_forward_length u32              = 18944\n",
      "llama_model_loader: - kv  10:                 qwen2.attention.head_count u32              = 28\n",
      "llama_model_loader: - kv  11:              qwen2.attention.head_count_kv u32              = 4\n",
      "llama_model_loader: - kv  12:                       qwen2.rope.freq_base f32              = 10000.000000\n",
      "llama_model_loader: - kv  13:     qwen2.attention.layer_norm_rms_epsilon f32              = 0.000001\n",
      "llama_model_loader: - kv  14:                       tokenizer.ggml.model str              = gpt2\n",
      "llama_model_loader: - kv  15:                         tokenizer.ggml.pre str              = deepseek-r1-qwen\n",
      "llama_model_loader: - kv  16:                      tokenizer.ggml.tokens arr[str,152064]  = [\"!\", \"\\\"\", \"#\", \"$\", \"%\", \"&\", \"'\", ...\n",
      "llama_model_loader: - kv  17:                  tokenizer.ggml.token_type arr[i32,152064]  = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...\n",
      "llama_model_loader: - kv  18:                      tokenizer.ggml.merges arr[str,151387]  = [\"Ġ Ġ\", \"ĠĠ ĠĠ\", \"i n\", \"Ġ t\",...\n",
      "llama_model_loader: - kv  19:                tokenizer.ggml.bos_token_id u32              = 151646\n",
      "llama_model_loader: - kv  20:                tokenizer.ggml.eos_token_id u32              = 151643\n",
      "llama_model_loader: - kv  21:            tokenizer.ggml.padding_token_id u32              = 151643\n",
      "llama_model_loader: - kv  22:               tokenizer.ggml.add_bos_token bool             = true\n",
      "llama_model_loader: - kv  23:               tokenizer.ggml.add_eos_token bool             = false\n",
      "llama_model_loader: - kv  24:                    tokenizer.chat_template str              = {% if not add_generation_prompt is de...\n",
      "llama_model_loader: - kv  25:               general.quantization_version u32              = 2\n",
      "llama_model_loader: - kv  26:                          general.file_type u32              = 15\n",
      "llama_model_loader: - type  f32:  141 tensors\n",
      "llama_model_loader: - type q4_K:  169 tensors\n",
      "llama_model_loader: - type q6_K:   29 tensors\n",
      "print_info: file format = GGUF V3 (latest)\n",
      "print_info: file type   = Q4_K - Medium\n",
      "print_info: file size   = 4.36 GiB (4.91 BPW) \n",
      "init_tokenizer: initializing tokenizer for type 2\n",
      "load: control token: 151660 '<|fim_middle|>' is not marked as EOG\n",
      "load: control token: 151659 '<|fim_prefix|>' is not marked as EOG\n",
      "load: control token: 151653 '<|vision_end|>' is not marked as EOG\n",
      "load: control token: 151645 '<｜Assistant｜>' is not marked as EOG\n",
      "load: control token: 151644 '<｜User｜>' is not marked as EOG\n",
      "load: control token: 151655 '<|image_pad|>' is not marked as EOG\n",
      "load: control token: 151651 '<|quad_end|>' is not marked as EOG\n",
      "load: control token: 151646 '<｜begin▁of▁sentence｜>' is not marked as EOG\n",
      "load: control token: 151643 '<｜end▁of▁sentence｜>' is not marked as EOG\n",
      "load: control token: 151652 '<|vision_start|>' is not marked as EOG\n",
      "load: control token: 151647 '<|EOT|>' is not marked as EOG\n",
      "load: control token: 151654 '<|vision_pad|>' is not marked as EOG\n",
      "load: control token: 151656 '<|video_pad|>' is not marked as EOG\n",
      "load: control token: 151661 '<|fim_suffix|>' is not marked as EOG\n",
      "load: control token: 151650 '<|quad_start|>' is not marked as EOG\n",
      "load: special_eos_id is not in special_eog_ids - the tokenizer config may be incorrect\n",
      "load: special tokens cache size = 22\n",
      "load: token to piece cache size = 0.9310 MB\n",
      "print_info: arch             = qwen2\n",
      "print_info: vocab_only       = 0\n",
      "print_info: n_ctx_train      = 131072\n",
      "print_info: n_embd           = 3584\n",
      "print_info: n_layer          = 28\n",
      "print_info: n_head           = 28\n",
      "print_info: n_head_kv        = 4\n",
      "print_info: n_rot            = 128\n",
      "print_info: n_swa            = 0\n",
      "print_info: n_embd_head_k    = 128\n",
      "print_info: n_embd_head_v    = 128\n",
      "print_info: n_gqa            = 7\n",
      "print_info: n_embd_k_gqa     = 512\n",
      "print_info: n_embd_v_gqa     = 512\n",
      "print_info: f_norm_eps       = 0.0e+00\n",
      "print_info: f_norm_rms_eps   = 1.0e-06\n",
      "print_info: f_clamp_kqv      = 0.0e+00\n",
      "print_info: f_max_alibi_bias = 0.0e+00\n",
      "print_info: f_logit_scale    = 0.0e+00\n",
      "print_info: f_attn_scale     = 0.0e+00\n",
      "print_info: n_ff             = 18944\n",
      "print_info: n_expert         = 0\n",
      "print_info: n_expert_used    = 0\n",
      "print_info: causal attn      = 1\n",
      "print_info: pooling type     = 0\n",
      "print_info: rope type        = 2\n",
      "print_info: rope scaling     = linear\n",
      "print_info: freq_base_train  = 10000.0\n",
      "print_info: freq_scale_train = 1\n",
      "print_info: n_ctx_orig_yarn  = 131072\n",
      "print_info: rope_finetuned   = unknown\n",
      "print_info: ssm_d_conv       = 0\n",
      "print_info: ssm_d_inner      = 0\n",
      "print_info: ssm_d_state      = 0\n",
      "print_info: ssm_dt_rank      = 0\n",
      "print_info: ssm_dt_b_c_rms   = 0\n",
      "print_info: model type       = 7B\n",
      "print_info: model params     = 7.62 B\n",
      "print_info: general.name     = DeepSeek R1 Distill Qwen 7B\n",
      "print_info: vocab type       = BPE\n",
      "print_info: n_vocab          = 152064\n",
      "print_info: n_merges         = 151387\n",
      "print_info: BOS token        = 151646 '<｜begin▁of▁sentence｜>'\n",
      "print_info: EOS token        = 151643 '<｜end▁of▁sentence｜>'\n",
      "print_info: EOT token        = 151643 '<｜end▁of▁sentence｜>'\n",
      "print_info: PAD token        = 151643 '<｜end▁of▁sentence｜>'\n",
      "print_info: LF token         = 198 'Ċ'\n",
      "print_info: FIM PRE token    = 151659 '<|fim_prefix|>'\n",
      "print_info: FIM SUF token    = 151661 '<|fim_suffix|>'\n",
      "print_info: FIM MID token    = 151660 '<|fim_middle|>'\n",
      "print_info: FIM PAD token    = 151662 '<|fim_pad|>'\n",
      "print_info: FIM REP token    = 151663 '<|repo_name|>'\n",
      "print_info: FIM SEP token    = 151664 '<|file_sep|>'\n",
      "print_info: EOG token        = 151643 '<｜end▁of▁sentence｜>'\n",
      "print_info: EOG token        = 151662 '<|fim_pad|>'\n",
      "print_info: EOG token        = 151663 '<|repo_name|>'\n",
      "print_info: EOG token        = 151664 '<|file_sep|>'\n",
      "print_info: max token length = 256\n",
      "load_tensors: loading model tensors, this can take a while... (mmap = true)\n",
      "load_tensors: layer   0 assigned to device CPU\n",
      "load_tensors: layer   1 assigned to device CPU\n",
      "load_tensors: layer   2 assigned to device CPU\n",
      "load_tensors: layer   3 assigned to device CPU\n",
      "load_tensors: layer   4 assigned to device CPU\n",
      "load_tensors: layer   5 assigned to device CPU\n",
      "load_tensors: layer   6 assigned to device CPU\n",
      "load_tensors: layer   7 assigned to device CPU\n",
      "load_tensors: layer   8 assigned to device CPU\n",
      "load_tensors: layer   9 assigned to device CPU\n",
      "load_tensors: layer  10 assigned to device CPU\n",
      "load_tensors: layer  11 assigned to device CPU\n",
      "load_tensors: layer  12 assigned to device CPU\n",
      "load_tensors: layer  13 assigned to device CPU\n",
      "load_tensors: layer  14 assigned to device CPU\n",
      "load_tensors: layer  15 assigned to device CPU\n",
      "load_tensors: layer  16 assigned to device CPU\n",
      "load_tensors: layer  17 assigned to device CPU\n",
      "load_tensors: layer  18 assigned to device CPU\n",
      "load_tensors: layer  19 assigned to device CPU\n",
      "load_tensors: layer  20 assigned to device CPU\n",
      "load_tensors: layer  21 assigned to device CPU\n",
      "load_tensors: layer  22 assigned to device CPU\n",
      "load_tensors: layer  23 assigned to device CPU\n",
      "load_tensors: layer  24 assigned to device CPU\n",
      "load_tensors: layer  25 assigned to device CPU\n",
      "load_tensors: layer  26 assigned to device CPU\n",
      "load_tensors: layer  27 assigned to device CPU\n",
      "load_tensors: layer  28 assigned to device CPU\n",
      "load_tensors: tensor 'token_embd.weight' (q4_K) (and 338 others) cannot be used with preferred buffer type CPU_AARCH64, using CPU instead\n",
      "load_tensors:   CPU_Mapped model buffer size =  4460.45 MiB\n",
      "....................................................................................\n",
      "llama_init_from_model: n_seq_max     = 1\n",
      "llama_init_from_model: n_ctx         = 512\n",
      "llama_init_from_model: n_ctx_per_seq = 512\n",
      "llama_init_from_model: n_batch       = 512\n",
      "llama_init_from_model: n_ubatch      = 512\n",
      "llama_init_from_model: flash_attn    = 0\n",
      "llama_init_from_model: freq_base     = 10000.0\n",
      "llama_init_from_model: freq_scale    = 1\n",
      "llama_init_from_model: n_ctx_per_seq (512) < n_ctx_train (131072) -- the full capacity of the model will not be utilized\n",
      "llama_kv_cache_init: kv_size = 512, offload = 1, type_k = 'f16', type_v = 'f16', n_layer = 28, can_shift = 1\n",
      "llama_kv_cache_init: layer 0: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 1: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 2: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 3: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 4: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 5: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 6: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 7: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 8: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 9: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 10: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 11: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 12: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 13: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 14: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 15: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 16: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 17: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 18: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 19: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 20: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 21: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 22: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 23: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 24: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 25: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 26: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init: layer 27: n_embd_k_gqa = 512, n_embd_v_gqa = 512\n",
      "llama_kv_cache_init:        CPU KV buffer size =    28.00 MiB\n",
      "llama_init_from_model: KV self size  =   28.00 MiB, K (f16):   14.00 MiB, V (f16):   14.00 MiB\n",
      "llama_init_from_model:        CPU  output buffer size =     0.58 MiB\n",
      "llama_init_from_model:        CPU compute buffer size =   304.00 MiB\n",
      "llama_init_from_model: graph nodes  = 986\n",
      "llama_init_from_model: graph splits = 1\n",
      "CPU : SSE3 = 1 | SSSE3 = 1 | AVX = 1 | AVX2 = 1 | F16C = 1 | FMA = 1 | BMI2 = 1 | AVX512 = 1 | AVX512_VBMI = 1 | AVX512_VNNI = 1 | AVX512_BF16 = 1 | LLAMAFILE = 1 | OPENMP = 1 | AARCH64_REPACK = 1 | \n",
      "Model metadata: {'general.file_type': '15', 'tokenizer.ggml.add_eos_token': 'false', 'tokenizer.ggml.add_bos_token': 'true', 'tokenizer.ggml.bos_token_id': '151646', 'qwen2.attention.layer_norm_rms_epsilon': '0.000001', 'general.architecture': 'qwen2', 'tokenizer.ggml.padding_token_id': '151643', 'general.basename': 'DeepSeek-R1-Distill-Qwen', 'qwen2.embedding_length': '3584', 'tokenizer.ggml.pre': 'deepseek-r1-qwen', 'general.name': 'DeepSeek R1 Distill Qwen 7B', 'qwen2.block_count': '28', 'general.type': 'model', 'general.size_label': '7B', 'general.license': 'mit', 'qwen2.context_length': '131072', 'tokenizer.chat_template': \"{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\\\n' + '```json' + '\\\\n' + tool['function']['arguments'] + '\\\\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\\\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\\\n' + '```json' + '\\\\n' + tool['function']['arguments'] + '\\\\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{{'<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\\\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}\", 'qwen2.attention.head_count_kv': '4', 'general.quantization_version': '2', 'tokenizer.ggml.model': 'gpt2', 'qwen2.feed_forward_length': '18944', 'qwen2.attention.head_count': '28', 'tokenizer.ggml.eos_token_id': '151643', 'qwen2.rope.freq_base': '10000.000000'}\n",
      "Available chat formats from metadata: chat_template.default\n",
      "Using gguf chat template: {% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{{'<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}\n",
      "Using chat eos_token: <｜end▁of▁sentence｜>\n",
      "Using chat bos_token: <｜begin▁of▁sentence｜>\n"
     ]
    }
   ],
   "source": [
    "# Create the provider.\n",
    "llama_model = Llama(\"./deepseek-r1-distill-qwen-7b-q4_k_m.gguf\", n_batch=1024, n_threads=10, n_gpu_layers=40)\n",
    "provider = LlamaCppPythonProvider(llama_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab2779f2-b483-4ddf-968c-818a25651a58",
   "metadata": {},
   "source": [
    "## 地磁模型测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "47da1b58-bc98-49a0-8188-53e6ea5a542e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_cpp_agent.llm_output_settings import LlmStructuredOutputSettings\n",
    "from typing import Tuple\n",
    "from llama_cpp_agent import MessagesFormatterType\n",
    "import math\n",
    "from datetime import datetime\n",
    "from pydantic import BaseModel\n",
    "import requests\n",
    "import json\n",
    "# 定义参数模型\n",
    "class ValueWithUnit(BaseModel):\n",
    "    value:float\n",
    "    unit:str\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "71abaa04-ab81-43c3-a548-917aad589bf3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "Read and follow the instructions below:\n",
      "\n",
      "<system_instructions>\n",
      "你是一个钻井智能体，只回答问题，不显示思考过程\n",
      "</system_instructions>\n",
      "\n",
      "\n",
      "You can call functions to help you with your tasks and user queries. The available functions are:\n",
      "\n",
      "<function_list>\n",
      "Function: calc_geomag\n",
      "  Description: 计算地磁信息\n",
      "  Parameters:\n",
      "    geo_type (str): 地磁模型\n",
      "    sample_time (str): 采样时间\n",
      "    alt (ValueWithUnit): 海拔高度\n",
      "    lat (float): 纬度\n",
      "    lon (float): 经度\n",
      "\n",
      "Model: ValueWithUnit\n",
      "  Fields:\n",
      "    value (float)\n",
      "    unit (str)\n",
      "</function_list>\n",
      "\n",
      "To call a function, respond with a JSON object (to call one function) or a list of JSON objects (to call multiple functions), with each object containing these fields:\n",
      "\n",
      "- \"function\": Put the name of the function to call here. \n",
      "- \"arguments\": Put the arguments to pass to the function here.\n",
      "\n",
      "The result of each function call will be returned to you before you need to respond again.<|im_end|>\n",
      "<|im_start|>user\n",
      "使用WMM模型，采样时间2024年10月31日，北纬21度31分，东经116度56分36秒, 海拔100米的地磁信息<|im_end|>\n",
      "<|im_start|>assistant"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_perf_context_print:        load time =    7661.99 ms\n",
      "llama_perf_context_print: prompt eval time =    7661.73 ms /   318 tokens (   24.09 ms per token,    41.51 tokens per second)\n",
      "llama_perf_context_print:        eval time =   13356.05 ms /    99 runs   (  134.91 ms per token,     7.41 tokens per second)\n",
      "llama_perf_context_print:       total time =   22072.00 ms /   417 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[\n",
      "\n",
      "{\n",
      "  \"function\":     \"calc_geomag\",\n",
      "  \"arguments\": {\n",
      "    \"geo_type\": \"WMM\",\n",
      "    \"sample_time\": \"2024-10-31\",\n",
      "    \"alt\": {\n",
      "      \"value\": 100,\n",
      "      \"unit\": \"meters\"\n",
      "    },\n",
      "    \"lat\": 21.5167,\n",
      "    \"lon\": 116.944\n",
      "  }\n",
      "}\n",
      "]\n",
      "[{'function': 'calc_geomag', 'arguments': {'geo_type': 'WMM', 'sample_time': '2024-10-31', 'alt': {'value': 100, 'unit': 'meters'}, 'lat': 21.5167, 'lon': 116.944}, 'return_value': {'detail': [{'loc': ['query', 'alt_u'], 'msg': \"value is not a valid enumeration member; permitted: 'm', 'meter', 'ft', 'feet'\", 'type': 'type_error.enum', 'ctx': {'enum_values': ['m', 'meter', 'ft', 'feet']}}]}}]\n"
     ]
    }
   ],
   "source": [
    "# Lets define a simple function tool\n",
    "def calc_geomag(geo_type:str, sample_time:str, alt:ValueWithUnit, lat:float, lon:float):\n",
    "    \"\"\"\n",
    "    计算地磁信息\n",
    "\n",
    "    Args:\n",
    "        geo_type: 地磁模型\n",
    "        sample_time: 采样时间\n",
    "        alt: 海拔高度\n",
    "        lat: 纬度\n",
    "        lon: 经度\n",
    "    \"\"\"\n",
    "    url = f'http://192.168.1.37:5116/geomag?mag_type={geo_type}&alt_v={alt.value}&alt_u={alt.unit}&dt={sample_time}&lon={lon}&lat={lat}'\n",
    "    r = requests.get(url)\n",
    "\n",
    "    return json.loads(r.content)\n",
    "\n",
    "# Now let's create an instance of the LlmStructuredOutput class by calling the `from_functions` function of it and passing it a list of functions.\n",
    "\n",
    "output_settings = LlmStructuredOutputSettings.from_functions([calc_geomag], allow_parallel_function_calling=True)\n",
    "\n",
    "# Create a LlamaCppAgent instance as before, including a system message with information about the tools available for the LLM agent.\n",
    "llama_cpp_agent = LlamaCppAgent(\n",
    "    provider,\n",
    "    debug_output=True,\n",
    "    system_prompt=f\"你是一个钻井智能体，只回答问题，不显示思考过程\",\n",
    "    predefined_messages_formatter_type=MessagesFormatterType.CHATML,\n",
    ")\n",
    "\n",
    "# Define some user input\n",
    "user_input = \"使用WMM模型，采样时间2024年10月31日，北纬21度31分，东经116度56分36秒, 海拔100米的地磁信息\"\n",
    "\n",
    "# Pass the user input together with output settings to `get_chat_response` method.\n",
    "# This will print the result of the function the LLM will call, it is a list of dictionaries containing the result.\n",
    "print(\n",
    "    llama_cpp_agent.get_chat_response(\n",
    "        user_input,\n",
    "        structured_output_settings=output_settings\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c850a52f-f82f-40ed-aa60-0467b8eace90",
   "metadata": {},
   "source": [
    "## 钻前设计测试"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44b769a-6170-487c-848a-320700530a4b",
   "metadata": {},
   "source": [
    "### 三段式测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "9ab89149-0858-4faf-8bbd-bf1e8ced3fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DesignOption(BaseModel):\n",
    "    build1:ValueWithUnit=Field(description=\"造斜率1\")\n",
    "    cl1:ValueWithUnit=Field(description=\"段长1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9934e84a-70a2-44e3-9d03-0df25f2cddb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Llama.generate: 109 prefix-match hit, remaining 192 prompt tokens to eval\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "Read and follow the instructions below:\n",
      "\n",
      "<system_instructions>\n",
      "你是一个钻井智能体，只回答问题，不显示思考过程\n",
      "</system_instructions>\n",
      "\n",
      "\n",
      "You can call functions to help you with your tasks and user queries. The available functions are:\n",
      "\n",
      "<function_list>\n",
      "Function: slant\n",
      "  Description: 三段式\n",
      "  Parameters:\n",
      "    opts (DesignOption): 设计选项\n",
      "\n",
      "Model: DesignOption\n",
      "  Fields:\n",
      "    build1 (ValueWithUnit): 造斜率1\n",
      "    cl1 (ValueWithUnit): 段长1\n",
      "\n",
      "Model: ValueWithUnit\n",
      "  Fields:\n",
      "    value (float)\n",
      "    unit (str)\n",
      "\n",
      "Model: ValueWithUnit\n",
      "  Fields:\n",
      "    value (float)\n",
      "    unit (str)\n",
      "</function_list>\n",
      "\n",
      "To call a function, respond with a JSON object (to call one function) or a list of JSON objects (to call multiple functions), with each object containing these fields:\n",
      "\n",
      "- \"function\": Put the name of the function to call here. \n",
      "- \"arguments\": Put the arguments to pass to the function here.\n",
      "\n",
      "The result of each function call will be returned to you before you need to respond again.<|im_end|>\n",
      "<|im_start|>user\n",
      "使用三段式,其中第一段段长100米，第一段造斜率为3度每30米<|im_end|>\n",
      "<|im_start|>assistant"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_perf_context_print:        load time =    7661.99 ms\n",
      "llama_perf_context_print: prompt eval time =    4588.01 ms /   192 tokens (   23.90 ms per token,    41.85 tokens per second)\n",
      "llama_perf_context_print:        eval time =   10547.11 ms /    78 runs   (  135.22 ms per token,     7.40 tokens per second)\n",
      "llama_perf_context_print:       total time =   15937.28 ms /   270 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[\n",
      "  {\n",
      "    \"function\": \n",
      "    \"slant\",\n",
      "    \"arguments\": {\n",
      "      \"opts\": {\n",
      "        \"build1\": {\n",
      "          \"value\": 3,\n",
      "          \"unit\": \"degree\"\n",
      "        },\n",
      "        \"cl1\": {\n",
      "          \"value\": 100,\n",
      "          \"unit\": \"meter\"\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "]\n",
      "[{'function': 'slant', 'arguments': {'opts': {'build1': {'value': 3, 'unit': 'degree'}, 'cl1': {'value': 100, 'unit': 'meter'}}}, 'return_value': {'detail': [{'loc': ['body', 'options'], 'msg': 'field required', 'type': 'value_error.missing'}, {'loc': ['body', 'target'], 'msg': 'field required', 'type': 'value_error.missing'}, {'loc': ['body', 'datum'], 'msg': 'field required', 'type': 'value_error.missing'}]}}]\n"
     ]
    }
   ],
   "source": [
    "# Lets define a simple function tool\n",
    "def slant(opts:DesignOption):\n",
    "    \"\"\"\n",
    "    三段式\n",
    "    \n",
    "    Args:\n",
    "        opts: 设计选项\n",
    "    \"\"\"\n",
    "    js_data = {\n",
    "      \"options\": {\n",
    "        \"checked\": \"build1,cl1\",\n",
    "        \"build1\": opts.build1,\n",
    "        \"cl1\":opts.cl1, \n",
    "        \"max_angle\": [\n",
    "          0,\n",
    "          \"deg\"\n",
    "        ],\n",
    "        \"cl2\": [\n",
    "          0,\n",
    "          \"meter\"\n",
    "        ]\n",
    "      },\n",
    "      \"target\": {\n",
    "        \"md\": [\n",
    "          0,\n",
    "          \"meter\"\n",
    "        ],\n",
    "        \"inc\": [\n",
    "          0,\n",
    "          \"deg\"\n",
    "        ],\n",
    "        \"azi\": [\n",
    "          0,\n",
    "          \"deg\"\n",
    "        ],\n",
    "        \"tvd\": [\n",
    "          0,\n",
    "          \"meter\"\n",
    "        ],\n",
    "        \"ns\": [\n",
    "          0,\n",
    "          \"meter\"\n",
    "        ],\n",
    "        \"ew\": [\n",
    "          0,\n",
    "          \"meter\"\n",
    "        ]\n",
    "      },\n",
    "      \"datum\": {\n",
    "        \"datum_elevation\": [\n",
    "          0,\n",
    "          \"meter\"\n",
    "        ],\n",
    "        \"convergence\": [\n",
    "          0,\n",
    "          \"deg\"\n",
    "        ]\n",
    "      }\n",
    "    }\n",
    "    url = f'http://192.168.1.37:5116/design/slant?id=1&api=false'\n",
    "    r = requests.post(url, data=js_data)\n",
    "    return json.loads(r.content)\n",
    "\n",
    "# Now let's create an instance of the LlmStructuredOutput class by calling the `from_functions` function of it and passing it a list of functions.\n",
    "\n",
    "output_settings = LlmStructuredOutputSettings.from_functions([slant], allow_parallel_function_calling=True)\n",
    "\n",
    "# Create a LlamaCppAgent instance as before, including a system message with information about the tools available for the LLM agent.\n",
    "llama_cpp_agent = LlamaCppAgent(\n",
    "    provider,\n",
    "    debug_output=True,\n",
    "    system_prompt=f\"你是一个钻井智能体，只回答问题，不显示思考过程\",\n",
    "    predefined_messages_formatter_type=MessagesFormatterType.CHATML,\n",
    ")\n",
    "\n",
    "# Define some user input\n",
    "user_input = \"使用三段式,其中第一段段长100米，第一段造斜率为3度每30米\"\n",
    "\n",
    "# Pass the user input together with output settings to `get_chat_response` method.\n",
    "# This will print the result of the function the LLM will call, it is a list of dictionaries containing the result.\n",
    "print(\n",
    "    llama_cpp_agent.get_chat_response(\n",
    "        user_input,\n",
    "        structured_output_settings=output_settings\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8263e2d8-162b-4b67-adaa-bd17f6f31a76",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
